import sys

import numpy as np
import pandas as pd
import wgdi.base as base

lens1 = base.newlens('vv.lens', 'order')
lens2 = base.newlens('pp.lens', 'order')
gff1 = base.newgff('vv160s.gff')
gff2 = base.newgff('pp72s.gff')
gff1 = gff1[gff1['chr'].isin(lens1.index)].copy()
gff2 = gff2[gff2['chr'].isin(lens2.index)].copy()

blast = base.newblast('vv160s_pp72s.blast', int(100),
                      float(1e-5), gff1, gff2, 'false')

blast['loc1'] = blast[0].map(gff1['order'])
blast['loc2'] = blast[1].map(gff2['order'])


df = pd.DataFrame(np.zeros((lens1.values[0], lens2.values[0])))
repeat_number = 5
rednum = 1
bluenum = 5+rednum
blast['color'] = 0
index = [group.sort_values(by=[11], ascending=[False])[:repeat_number].index.tolist()
                 for name, group in blast.groupby([0])]
reddata = np.array([k[:rednum] for k in index], dtype=object)
bluedata = np.array([k[rednum:bluenum] for k in index], dtype=object)
graydata = np.array([k[bluenum:repeat_number] for k in index], dtype=object)
if len(reddata):
    redindex = np.concatenate(reddata)
else:
    redindex = []
if len(bluedata):
    blueindex = np.concatenate(bluedata)
else:
    blueindex = []
if len(graydata):
    grayindex = np.concatenate(graydata)
else:
    grayindex = []
blast.loc[redindex, 'color'] = 50
blast.loc[blueindex, 'color'] = 40
blast.loc[grayindex, 'color'] = 25

for index, row in blast.iterrows():
    df.loc[row['loc1']-1, row['loc2']-1] = row['color']
print(df.head())
df = df.loc[:, df.sum(axis=0) != 0]
df = df.loc[df.sum(axis=1) != 0, :]
df.to_csv('matrix.csv')
